package algorithms.prefix;

/**
 * 前缀树
 */
public class PriFixTree {
    private Node root;

    public PriFixTree() {
        this.root = new Node();
    }

    class Node {
        private int pass;// 记录经过的次数
        private int end;// 记录结尾的次数
        private Node[] nodes;// 当前节点下的路径

        public Node() {
            pass = 0;
            end = 0;
            // a-z 有24个字符，0表示a,1表示b...
            this.nodes = new Node[24];
        }
    }

    /**
     * 插入前缀树
     *
     * @param word
     */
    public void insert(String word) {
        if (word == null) {
            return;
        }
        char[] chars = word.toCharArray();
        Node node = this.root;
        node.pass++;
        int index = 0;
        for (int i = 0; i < chars.length; i++) {
            index = chars[i] - 'a';
            if (node.nodes[index] == null) {
                node.nodes[index] = new Node();
            }
            node = node.nodes[index];// 向下一个节点走
            node.pass++;
        }
        // 到结尾
        node.end++;

    }

    /**
     * 查询字符word出现的次数
     *
     * @param word
     * @return 出现次数
     */
    public int search(String word) {
        if (word == null) {
            return 0;
        }
        char[] chars = word.toCharArray();
        Node node = this.root;
        int index = 0;
        for (int i = 0; i < chars.length; i++) {
            index = chars[i] - 'a';
            if (node.nodes[index] == null) {// 查询结果没走完，树走完了：不存在结果
               return 0;
            }
            node = node.nodes[index];// 向下一个节点走
        }
        // 到结尾
        return node.end;
    }

    /**
     * 查询以prefix为前缀的字符次数
     *
     * @param prefix
     * @return 出现次数
     */
    public int searchPrefix(String prefix) {
        if (prefix == null) {
            return 0;
        }
        char[] chars = prefix.toCharArray();
        Node node = this.root;
        int index = 0;
        for (int i = 0; i < chars.length; i++) {
            index = chars[i] - 'a';
            if (node.nodes[index] == null) {// 查询结果没走完，树走完了：不存在结果
                return 0;
            }
            node = node.nodes[index];// 向下一个节点走
        }
        // 到结尾
        return node.pass;
    }

    /**
     * 从前缀树中移除word
     *
     * @param word
     */
    public void remove(String word) {
        if (search(word) == 0) {// 存在才删除
            return;
        }
        char[] chars = word.toCharArray();
        Node node = this.root;
        node.pass--;
        int index = 0;
        for (int i = 0; i < chars.length; i++) {
            index = chars[i] - 'a';
            if (--node.nodes[index].pass == 0) {// pass 为0了后面的树可以全部丢掉
                node.nodes[index] = null;
                return;
            }
            node = node.nodes[index];// 向下一个节点走
        }
        // 到结尾
        node.end--;
    }

}
